﻿using System;
using System.Collections.Generic;
using System.Globalization;
using System.IO;
using System.Xml;
using CoreRay.Cameras;

namespace CoreRay.IO
{
    public class ColladaResourceHandler : ResourceHandler
    {
        public override string[] Extensions
        {
            get 
            {
                return new string[] { "dae" }; 
            }
        }

        public override bool CanRead
        {
            get 
            {
                return true;
            }
        }

        public override bool CanWrite
        {
            get 
            {
                return false;
            }
        }

        public override object Read(Stream input)
        {
            // Create the scene object
            Scene scene = new Scene();

            XmlReaderSettings settings = new XmlReaderSettings();
            settings.IgnoreWhitespace = true;
            settings.IgnoreComments = true;

            ColladaHeader colladaHeader = new ColladaHeader();
            Dictionary<string, Mesh> meshList = null;

            // Open the document
            using (XmlReader reader = XmlReader.Create(input, settings))
            {
                while (reader.Read())
                {
                    if (reader.NodeType == XmlNodeType.Element)
                    {
                        switch (reader.Name)
                        {
                            case "asset":
                                colladaHeader = ReadAssetHeader(reader);
                                break;

                            case "library_geometries":
                                meshList = ReadLibraryGeometries(scene, colladaHeader, reader);
                                break;

                            case "library_visual_scenes":
                                ReadSceneDescription(scene, colladaHeader, meshList, reader);
                                break;
                        }
                    }
                }
            }

            if (scene.MainCamera == null)
            {
                scene.MainCamera = new PerspectiveCamera(new Vector3(0.0f, 30.0f, -200.0f), new Vector3(0, 20, 0));
                scene.AddSceneObject(scene.MainCamera);
            }

            // Return the scene object
            return scene;
        }

        private ColladaHeader ReadAssetHeader(XmlReader reader)
        {
            var colladaHeader = new ColladaHeader();

            while (reader.Read())
            {
                if (reader.NodeType == XmlNodeType.EndElement && reader.Name == "asset")
                {
                    break;
                }

                else if (reader.NodeType == XmlNodeType.Element && reader.Name == "unit")
                {
                    reader.MoveToAttribute("meter");
                    colladaHeader.MeterConversionRatio = float.Parse(reader.Value, CultureInfo.InvariantCulture);
                }

                else if (reader.NodeType == XmlNodeType.Element && reader.Name == "up_axis")
                {
                    reader.Read();

                    switch (reader.Value)
                    {
                        case "X_UP":
                            colladaHeader.SceneOrientation = ColladaSceneOrientation.UpX;
                            break;

                        case "Y_UP":
                            colladaHeader.SceneOrientation = ColladaSceneOrientation.UpY;
                            break;

                        case "Z_UP":
                            colladaHeader.SceneOrientation = ColladaSceneOrientation.UpZ;
                            break;
                    }
                }
            }

            return colladaHeader;
        }

        private void ReadSceneDescription(Scene scene, ColladaHeader colladaHeader, Dictionary<string, Mesh> meshList, XmlReader reader)
        {
            while (reader.Read())
            {
                if (reader.NodeType == XmlNodeType.EndElement && reader.Name == "library_visual_scenes")
                {
                    break;
                }

                else if (reader.NodeType == XmlNodeType.Element && reader.Name == "node")
                {
                    reader.MoveToAttribute("name");
                    var nodeName = reader.Value;

                    List<Matrix> transformations = new List<Matrix>();

                    while (reader.Read())
                    {
                        if (reader.NodeType == XmlNodeType.EndElement && reader.Name == "node")
                        {
                            break;
                        }

                        else if (reader.NodeType == XmlNodeType.Element && reader.Name == "node")
                        {
                            reader.MoveToAttribute("name");
                            nodeName = reader.Value;
                        }

                        else if (reader.NodeType == XmlNodeType.Element && reader.Name == "translate")
                        {
                            reader.Read();
                            string[] translateElements = reader.Value.Split(new char[] { ' ' }, StringSplitOptions.RemoveEmptyEntries);

                            if (colladaHeader.SceneOrientation == ColladaSceneOrientation.UpZ)
                            {
                                transformations.Add(Matrix.Translation(new Vector3(float.Parse(translateElements[0], CultureInfo.InvariantCulture), float.Parse(translateElements[2], CultureInfo.InvariantCulture), float.Parse(translateElements[1], CultureInfo.InvariantCulture))));
                            }

                            else
                            {
                                transformations.Add(Matrix.Translation(new Vector3(float.Parse(translateElements[0], CultureInfo.InvariantCulture), float.Parse(translateElements[1], CultureInfo.InvariantCulture), -float.Parse(translateElements[2], CultureInfo.InvariantCulture))));
                            }
                        }

                        else if (reader.NodeType == XmlNodeType.Element && reader.Name == "scale")
                        {
                            reader.Read();
                            string[] scaleElements = reader.Value.Split(new char[] { ' ' }, StringSplitOptions.RemoveEmptyEntries);

                            if (colladaHeader.SceneOrientation == ColladaSceneOrientation.UpZ)
                            {
                                transformations.Add(Matrix.ScaleX(float.Parse(scaleElements[0])));
                                transformations.Add(Matrix.ScaleY(float.Parse(scaleElements[2])));
                                transformations.Add(Matrix.ScaleZ(float.Parse(scaleElements[1])));
                            }

                            else
                            {
                                transformations.Add(Matrix.ScaleX(float.Parse(scaleElements[0])));
                                transformations.Add(Matrix.ScaleY(float.Parse(scaleElements[1])));
                                transformations.Add(Matrix.ScaleZ(-float.Parse(scaleElements[2])));
                            }
                        }

                        else if (reader.NodeType == XmlNodeType.Element && reader.Name == "rotate")
                        {
                            reader.Read();
                            string[] rotateElements = reader.Value.Split(new char[] { ' ' }, StringSplitOptions.RemoveEmptyEntries);

                            float rotateAmount = float.Parse(rotateElements[3], CultureInfo.InvariantCulture);

                            if (rotateElements[0] != "0")
                            {
                                float rotateRatio = float.Parse(rotateElements[0], CultureInfo.InvariantCulture);

                                if (colladaHeader.SceneOrientation == ColladaSceneOrientation.UpZ)
                                {
                                    transformations.Add(Matrix.RotationX(MathUtility.DegreesToRadians(rotateRatio * -rotateAmount)));
                                }

                                else
                                {
                                    transformations.Add(Matrix.RotationX(MathUtility.DegreesToRadians(rotateRatio * -rotateAmount)));
                                }
                            }

                            if (rotateElements[2] != "0")
                            {
                                float rotateRatio = float.Parse(rotateElements[2], CultureInfo.InvariantCulture);

                                if (colladaHeader.SceneOrientation == ColladaSceneOrientation.UpZ)
                                {
                                    transformations.Add(Matrix.RotationY(MathUtility.DegreesToRadians(rotateRatio * -rotateAmount)));
                                }

                                else
                                {
                                    transformations.Add(Matrix.RotationZ(MathUtility.DegreesToRadians(rotateAmount)));
                                }
                            }

                            if (rotateElements[1] != "0")
                            {
                                float rotateRatio = float.Parse(rotateElements[1], CultureInfo.InvariantCulture);

                                if (colladaHeader.SceneOrientation == ColladaSceneOrientation.UpZ)
                                {
                                    transformations.Add(Matrix.RotationZ(MathUtility.DegreesToRadians(rotateRatio * -rotateAmount)));
                                }

                                else
                                {
                                    transformations.Add(Matrix.RotationY(MathUtility.DegreesToRadians(-rotateAmount)));
                                }
                            }
                        }

                        else if (reader.NodeType == XmlNodeType.Element && reader.Name == "matrix")
                        {
                            reader.Read();
                            string[] matrixElements = reader.Value.Split(new char[] { ' ' }, StringSplitOptions.RemoveEmptyEntries);

                            Matrix matrix = new Matrix(float.Parse(matrixElements[0], CultureInfo.InvariantCulture), float.Parse(matrixElements[4], CultureInfo.InvariantCulture), float.Parse(matrixElements[8], CultureInfo.InvariantCulture), float.Parse(matrixElements[12], CultureInfo.InvariantCulture),
                                                        float.Parse(matrixElements[1], CultureInfo.InvariantCulture), float.Parse(matrixElements[5], CultureInfo.InvariantCulture), float.Parse(matrixElements[9], CultureInfo.InvariantCulture), float.Parse(matrixElements[13], CultureInfo.InvariantCulture),
                                                        float.Parse(matrixElements[2], CultureInfo.InvariantCulture), float.Parse(matrixElements[6], CultureInfo.InvariantCulture), float.Parse(matrixElements[10], CultureInfo.InvariantCulture), float.Parse(matrixElements[14], CultureInfo.InvariantCulture),
                                                        float.Parse(matrixElements[3], CultureInfo.InvariantCulture), float.Parse(matrixElements[7], CultureInfo.InvariantCulture), float.Parse(matrixElements[11], CultureInfo.InvariantCulture), float.Parse(matrixElements[15], CultureInfo.InvariantCulture));

                            if (colladaHeader.SceneOrientation == ColladaSceneOrientation.UpZ)
                            {
                                Matrix convertMatrix = new Matrix(1, 0, 0, 0,
                                                                  0, 0, 1, 0,
                                                                  0, 1, 0, 0,
                                                                  0, 0, 0, 1);

                                matrix = matrix * convertMatrix;
                            }

                            transformations.Add(matrix);
                        }

                        else if (reader.NodeType == XmlNodeType.Element && reader.Name == "instance_geometry")
                        {
                            // Compute world matrix
                            Matrix worldMatrix = Matrix.Identity;

                            // All the transformations are done in reverse order because collada is right-handed
                            for (int i = transformations.Count - 1; i >= 0; i--)
                            {
                                worldMatrix = worldMatrix * transformations[i];
                            }

                            reader.MoveToAttribute("url");
                            var meshId = reader.Value.Replace("#", "");

                            if (meshList.ContainsKey(meshId))
                            {
                                var mesh = meshList[meshId];
                                mesh.Name = nodeName;
                                mesh.WorldMatrix = worldMatrix;
                                scene.AddSceneObject(mesh);
                            }
                        }

                        else if (reader.NodeType == XmlNodeType.Element && reader.Name == "instance_camera")
                        {
                            // Compute world matrix
                            Matrix worldMatrix = Matrix.ScaleY(-1) * Matrix.RotationX(MathUtility.DegreesToRadians(180));

                            for (int i = transformations.Count - 1; i >= 0; i--)
                            {
                                worldMatrix = worldMatrix * transformations[i];
                            }

                            var camera = new PerspectiveCamera(worldMatrix.Inverse());
                            camera.Name = nodeName;
                            scene.AddSceneObject(camera);

                            if (scene.MainCamera == null)
                            {
                                scene.MainCamera = camera;
                            }
                        }
                    }
                }
            }
        }

        private Dictionary<string, Mesh> ReadLibraryGeometries(Scene scene, ColladaHeader colladaHeader, XmlReader reader)
        {
            Dictionary<string, Mesh> meshList = new Dictionary<string, Mesh>();

            while (reader.Read())
            {
                if (reader.NodeType == XmlNodeType.EndElement && reader.Name == "library_geometries")
                {
                    break;
                }

                else if (reader.NodeType == XmlNodeType.Element && reader.Name == "geometry")
                {
                    var mesh = ReadGeometryObject(scene, colladaHeader, reader);
                    meshList.Add(mesh.Name, mesh);
                }
            }

            return meshList;
        }

        private Mesh ReadGeometryObject(Scene scene, ColladaHeader colladaHeader, XmlReader reader)
        {
            // Get the name of the object
            reader.MoveToAttribute("id");
            string objectName = reader.Value;

            while (reader.Read())
            {
                if (reader.NodeType == XmlNodeType.EndElement && reader.Name == "geometry")
                {
                    break;
                }

                else if (reader.NodeType == XmlNodeType.Element && reader.Name == "mesh")
                {
                    var mesh = ReadMeshObject(scene, colladaHeader, reader);
                    mesh.Name = objectName;
                    return mesh;
                }
            }

            return null;
        }

        private Mesh ReadMeshObject(Scene scene, ColladaHeader colladaHeader, XmlReader reader)
        {
            var mesh = new Mesh();
            Dictionary<string, ColladaSource> dataSources = new Dictionary<string, ColladaSource>();
            Dictionary<string, string> vertexElementDatasource = new Dictionary<string, string>();
            string vertexDataSourceName = null;
            string vertexDataSourceKey = null;
            List<int[]> elementIndices = new List<int[]>();

            while (reader.Read())
            {
                if (reader.NodeType == XmlNodeType.EndElement && reader.Name == "mesh")
                {
                    break;
                }

                else if (reader.NodeType == XmlNodeType.Element && reader.Name == "source")
                {
                    var dataSource = new ColladaSource();

                    // Get the id of the source
                    reader.MoveToAttribute("id");
                    dataSource.Id = reader.Value;

                    dataSources.Add(dataSource.Id, dataSource);

                    // Read the source according in the correct type
                    reader.Read();

                    switch (reader.Name)
                    {
                        case "float_array":

                            dataSource.SourceType = ColladaSourceType.Float;
                            dataSource.FloatData = ReadFloatArraySource(reader);
                            break;
                    } 

                    // Read data source layout
                    reader.Read();
                    reader.Read();
                    reader.Read();

                    // Read data source attributes
                    reader.MoveToAttribute("count");
                    dataSource.ElementCount = int.Parse(reader.Value, CultureInfo.InvariantCulture);

                    reader.MoveToAttribute("stride");
                    dataSource.ElementSize = int.Parse(reader.Value, CultureInfo.InvariantCulture);

                    int currentIndex = 0;

                    while (reader.Read())
                    {
                        if (reader.NodeType == XmlNodeType.EndElement && reader.Name == "accessor")
                        {
                            break;
                        }

                        else if (reader.NodeType == XmlNodeType.Element && reader.Name == "param")
                        {
                            reader.MoveToAttribute("name");
                            var name = reader.Value;
                            dataSource.ElementLocalIndexList.Add(name, currentIndex);
                            currentIndex++;
                        }
                    }
                }

                else if (reader.NodeType == XmlNodeType.Element && reader.Name == "vertices")
                {
                    reader.MoveToAttribute("id");
                    vertexDataSourceKey = reader.Value;

                    reader.Read();

                    reader.MoveToAttribute("source");
                    vertexDataSourceName = reader.Value.Replace("#", "");
                }

                else if (reader.NodeType == XmlNodeType.Element && (reader.Name == "polygons" || reader.Name == "triangles" || reader.Name == "polylist"))
                {
                    reader.MoveToAttribute("count");
                    var polygonCount = int.Parse(reader.Value, CultureInfo.InvariantCulture);

                    // Create the vertex format
                    VertexFormat vertexFormat = new VertexFormat();

                    while (reader.Read())
                    {
                        if (reader.NodeType == XmlNodeType.Element && reader.Name == "p")
                        {
                            break;
                        }

                        else if (reader.NodeType == XmlNodeType.Element && reader.Name == "input")
                        {
                            string elementName = null;

                            // Read semantic name
                            reader.MoveToAttribute("semantic");
                            var semanticName = reader.Value;

                            // Read datasource id
                            reader.MoveToAttribute("source");
                            var datasourceId = reader.Value.Replace("#", "");

                            switch (semanticName)
                            {
                                case "VERTEX":
                                    elementName = "Position";
                                    break;

                                case "NORMAL":
                                    elementName = "Normal";
                                    break;

                                case "TEXCOORD":
                                    elementName = "TextureCoordinates";
                                    break;
                            }

                            if (elementName != null)
                            {
                                Type elementType = typeof(Vector3);

                                var dataSource = (vertexDataSourceKey == datasourceId) ? dataSources[vertexDataSourceName] : dataSources[datasourceId];

                                if (dataSource.ElementSize == 2 || elementName == "TextureCoordinates")
                                {
                                    elementType = typeof(Vector2);
                                }

                                vertexElementDatasource.Add(elementName, dataSource.Id);
                                vertexFormat.AddElement(elementName, elementType);
                            }
                        }
                    }

                    // Create the geometry buffer
                    mesh.GeometryBuffer = new GeometryBuffer(vertexFormat);

                    // Process triangle indices
                    do
                    {
                        if (reader.NodeType == XmlNodeType.EndElement && (reader.Name == "polygons" || reader.Name == "triangles" || reader.Name == "polylist"))
                        {
                            break;
                        }

                        else if (reader.NodeType == XmlNodeType.Element && reader.Name == "p")
                        {
                            reader.Read();
                            string[] indices = reader.Value.Split(new char[] { ' ' }, StringSplitOptions.RemoveEmptyEntries);

                            for (int i = 0; i < indices.Length; i += vertexFormat.ElementsCount)
                            {
                                elementIndices.Add(new int[vertexFormat.ElementsCount]);

                                for (int j = 0; j < vertexFormat.ElementsCount; j++)
                                {
                                    elementIndices[elementIndices.Count - 1][j] = int.Parse(indices[i + j], CultureInfo.InvariantCulture);
                                } 
                            }     
                        }
                    } while (reader.Read());

                    // Reorder indices
                    for (int i = 0; i < elementIndices.Count; i++)
                    {
                        if ((i % 3) == 1)
                        {
                            int[] temp = elementIndices[i];
                            elementIndices[i] = elementIndices[i + 1];
                            elementIndices[i + 1] = temp;
                        }
                    }

                    Dictionary<string, int> vertexValueIndexMapping = new Dictionary<string, int>();
                    int indexCount = 0;

                    // Write the vertex elements
                    for (int i = 0; i < elementIndices.Count; i++)
                    {
                        string vertexSignature = string.Empty;
                        object[] valuesToWrite = new object[vertexFormat.ElementsCount];

                        for (int j=0;j<vertexFormat.ElementsCount;j++)
                        {
                            var vertexFormatElement = vertexFormat[j];

                            var dataSource = dataSources[vertexElementDatasource[vertexFormat[j].Name]];

                            if (dataSource.ElementSize == 3 && (vertexFormatElement.Name != "TextureCoordinates"))
                            {
                                Vector3 vector;

                                if (colladaHeader.SceneOrientation == ColladaSceneOrientation.UpZ)
                                {
                                    vector = new Vector3(dataSource.FloatData[elementIndices[i][j] * dataSource.ElementSize], dataSource.FloatData[elementIndices[i][j] * dataSource.ElementSize + 2], dataSource.FloatData[elementIndices[i][j] * dataSource.ElementSize + 1]);
                                }

                                else
                                {
                                    vector = new Vector3(dataSource.FloatData[elementIndices[i][j] * dataSource.ElementSize], dataSource.FloatData[elementIndices[i][j] * dataSource.ElementSize + 1], -dataSource.FloatData[elementIndices[i][j] * dataSource.ElementSize + 2]);
                                }

                                if(vertexFormat[j].Name == "Normal")
                                {
                                    vector = vector.Normalize();
                                }

                                valuesToWrite[j] = vector;
                                vertexSignature += vector.ToString();
                            }

                            else
                            {
                                var vector = new Vector2(dataSource.FloatData[elementIndices[i][j] * dataSource.ElementSize], dataSource.FloatData[elementIndices[i][j] * dataSource.ElementSize + 1]);

                                valuesToWrite[j] = vector;
                                vertexSignature += vector.ToString();
                            }
                        }

                        if (vertexValueIndexMapping.ContainsKey(vertexSignature))
                        {
                            mesh.GeometryBuffer.IndexBuffer.Write(vertexValueIndexMapping[vertexSignature]);
                        }

                        else
                        {
                            for (int j = 0; j < vertexFormat.ElementsCount; j++)
                            {
                                if (valuesToWrite[j] is Vector3)
                                {
                                    mesh.GeometryBuffer.VertexBuffer.Write((Vector3)valuesToWrite[j]);
                                }

                                else
                                {
                                    mesh.GeometryBuffer.VertexBuffer.Write((Vector2)valuesToWrite[j]);
                                }
                            }

                            vertexValueIndexMapping.Add(vertexSignature, indexCount);
                            mesh.GeometryBuffer.IndexBuffer.Write(indexCount);
                            indexCount++;
                        } 
                    } 
                }
            }

            mesh.Commit();
            return mesh;
        }

        private float[] ReadFloatArraySource(XmlReader reader)
        {
            // Get the id of the source
            reader.MoveToAttribute("count");
            var count = int.Parse(reader.Value, CultureInfo.InvariantCulture);

            // Get the data of the source
            reader.Read();
            string data = reader.Value;

            // Create the destination array
            float[] outputData = new float[count];

            // Parse the data source
            string[] dataElements = data.Replace('\n', ' ').Split(new char[] { ' ' }, StringSplitOptions.RemoveEmptyEntries);

            for (int i = 0; i < dataElements.Length; i++)
            {
                outputData[i] = float.Parse(dataElements[i], CultureInfo.InvariantCulture);
            }

            return outputData;
        }
    }
}
